import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, random_split
import numpy as np
import matplotlib.pyplot as plt

# 定义数据预处理（标准化）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 下载 MNIST 数据集
train_dataset = torchvision.datasets.MNIST(root="./data", train=True, transform=transform, download=True)
test_dataset = torchvision.datasets.MNIST(root="./data", train=False, transform=transform, download=True)

# 划分训练集和验证集
train_size = int(0.8 * len(train_dataset))  # 80% 训练集
val_size = len(train_dataset) - train_size  # 20% 验证集
train_dataset, val_dataset = random_split(train_dataset, [train_size, val_size])

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 测试数据
for images, labels in train_loader:
    print(f"图像批次大小: {images.shape}")  # (64, 1, 28, 28)
    print(f"标签批次大小: {labels.shape}")  # (64,)
    break